#pragma once

#include "ClamAVScanner.h"
#include "InfectedHandler.h"
#include "KernelHandler.h"
#include "Config.h"
#include "MD5Signature.h"
#include "InfectedDetails.h"
#include "BaseScanGroup.h"
#include "Thread.h"
#include "ThreadPool.h"
#include "CriticalSection.h"
#include "Semaphore.h"
#include <Windows.h>
#include <string>
#include <memory>
#include <vector>
#include <CL\cl.hpp>

namespace GPUCC
{
	class Scanner
	{
	public:
		class Engine : public InfectedHandler, KernelHandler
		{
		public:
			Engine(Scanner& Instance);
			virtual ~Engine();

			bool Start();
			void Stop();

			virtual void OnInfected(const std::string& Path, const std::string& SignatureName);
			virtual void OnKernelOutputRead(BaseScanGroup* Instance);
			virtual void OnGroupFinished(BaseScanGroup* Instance);
			virtual void OnStreamGroupFinished(BaseScanGroup* Instance);

		private:
			Scanner& m_Scanner;
			CPU::ThreadPool m_ThreadPool;
		};

		Scanner(const Config& Conf, bool CpuOnly = false);
		~Scanner();

		bool Init();
		void Cleanup();
		
		std::vector<InfectedDetails> ScanDrives(bool ScanRemovableDrives = false);
		std::vector<InfectedDetails> ScanFolder(std::string& Path, bool Recursive = true);

		unsigned int GetScannedFilesCount();

	private:
		bool InitGPU();
		bool InitOpenCL();
		bool LoadSources(std::vector<std::string>& strings, cl::Program::Sources& sources);
		bool StartThreads();
		void CompleteScan();

		void ScanSingleFolder(std::string& Path, bool Recursive = true);
		void ScanFile(const std::string& Path, const DWORD SizeLow, const DWORD SizeHigh);
		void ScanFile(const std::string& Path, const LARGE_INTEGER& Size);
		void ScanEntireFile(const std::string& Path, unsigned int FileSize);
		void ScanFileStreamed(const std::string& Path, unsigned int FileSize);

		bool m_CpuOnly;
		const Config& m_Config;
		std::vector<MD5Signature> m_Signatures;

		Engine m_Engine;
		unsigned int m_ScannedFilesCounter;

		std::vector<InfectedDetails> m_Results;
		CPU::CriticalSection m_ResultsCS;

		ClamAVScanner m_ClamAV;
		CPU::ThreadPool m_ClamAVThreads;

		cl::Platform m_Platform;
		cl::Device m_GPU;
		cl::Context m_Context;
		cl::Program m_Program;
		std::vector<cl::CommandQueue> m_CommandQueue;
		std::vector<cl::CommandQueue> m_StreamCommandQueue;

		std::vector<std::shared_ptr<BaseScanGroup>> m_Groups;
		std::vector<std::shared_ptr<BaseScanGroup>> m_StreamGroups;
		CPU::CriticalSection m_GroupsCS;
		CPU::CriticalSection m_StreamGroupsCS;
		CPU::Semaphore m_GroupsLimits;
		CPU::Semaphore m_StreamGroupsLimits;

		// todo: fix this!!!
		friend class ReportJob;
		friend class GroupFinishedJob;
		friend class StreamGroupFinishedJob;
	};
}